import torch
import time
import pynvml
import argparse


def occupy_gpu_memory(gpu_id, memory_fraction=0.8):
    """Occupy a specified fraction of GPU memory."""
    device = torch.device(f"cuda:{gpu_id}")
    torch.cuda.set_device(device)
    total_memory = torch.cuda.get_device_properties(device).total_memory
    rand = torch.rand(1) * 0.3 + 0.7
    block_memory = int(rand * total_memory * memory_fraction)

    # Allocate memory on the GPU
    _ = torch.empty(block_memory // 4, dtype=torch.float32, device=device)
    print(
        f"GPU {gpu_id} occupied with {memory_fraction * 100}% of total memory."
    )


def check_and_occupy_gpus(memory_fraction=0.8, interval=10, max_gpus=2):
    """Check for vacant GPUs and occupy them, but no more than max_gpus."""
    # Initialize NVML

    occupied_gpus = 0
    while True:
        pynvml.nvmlInit()
        # Get the number of GPUs
        gpu_count = pynvml.nvmlDeviceGetCount()
        # Iterate over all GPUs
        for gpu_id in range(gpu_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)

            # Get memory information
            memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)

            total_memory = (
                memory_info.total / 1024**2  # type: ignore
            )  # Convert bytes to MB
            used_memory = memory_info.used / 1024**2  # type: ignore
            # free_memory = memory_info.free / 1024**2  # type: ignore
            memory_utilization = (used_memory / total_memory) * 100
            if memory_utilization < 2 and occupied_gpus < max_gpus:  # type: ignore
                print(
                    f"GPU {gpu_id} is vacant {memory_utilization}. Trying to occupy it."
                )
                occupy_gpu_memory(gpu_id, memory_fraction)
                occupied_gpus += 1
        time.sleep(interval)
        # if occupied_gpus == max_gpus:
        #     time.sleep()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--memory_fraction", type=float, default=0.8)
    parser.add_argument("--interval", type=int, default=10)
    parser.add_argument("--max_gpus", type=int, default=2)
    args = parser.parse_args()
    check_and_occupy_gpus(
        memory_fraction=args.memory_fraction,
        interval=args.interval,
        max_gpus=args.max_gpus,
    )
